"""
Description: 阿里云证书申请工具
Author: 嘎嘣脆的贝爷
Date: 2024-11-05 17:29:28
LastEditTime: 2024-11-05 17:31:24
LastEditors: 嘎嘣脆的贝爷
"""

from alibabacloud_cas20200407.client import Client as cas20200407Client
from alibabacloud_cas20200407 import models as cas_20200407_models
from alibabacloud_tea_openapi import models as open_api_models
from alibabacloud_tea_util import models as util_models
from config import config
from tool import tool
import datetime, asyncio, math


class AliCertAgent:
    """
    阿里云证书申请工具
    """

    def __init__(self):
        # 要监测的域名列表
        self.domain_list = config.DOMAIN_LIST
        # 已完成列表
        self.finish_list = {}

        # 初始化客户端
        self.client = cas20200407Client(
            open_api_models.Config(
                access_key_id=config.ACCESS_KEY_ID,
                access_key_secret=config.ACCESS_KEY_SECRET,
                endpoint="cas.aliyuncs.com",
            )
        )

    def __get_cert_balance(self) -> int:
        """查询证书余额"""
        describe_package_state_request = (
            cas_20200407_models.DescribePackageStateRequest(
                product_code=config.PRODUCT_CODE
            )
        )
        runtime = util_models.RuntimeOptions()
        try:
            rst = self.client.describe_package_state_with_options(
                describe_package_state_request, runtime
            )
            return rst.body.total_count - rst.body.used_count
        except Exception as error:
            print("查询证书余额出错:", error.message)

    def __get_cert_list(self, status: str = None, order_type: str = "CPACK") -> list:
        """
        获取证书列表
        return: list[{}]
        """
        list_user_certificate_order_request = (
            cas_20200407_models.ListUserCertificateOrderRequest(
                order_type=order_type, status=status
            )
        )
        runtime = util_models.RuntimeOptions()
        try:
            rst = self.client.list_user_certificate_order_with_options(
                list_user_certificate_order_request, runtime
            )
            return rst.body.certificate_order_list
        except Exception as error:
            print("获取证书列表出错:", error.message)

    async def __create_cert(self, domain: str) -> dict:
        """
        创建证书
        :param domain: 域名
        :return: {订单id: 域名}
        """
        if not domain:
            raise ValueError("域名不能为空")
        create_certificate_for_package_request_request = (
            cas_20200407_models.CreateCertificateForPackageRequestRequest(
                product_code=config.PRODUCT_CODE, domain=domain, validate_type="DNS"
            )
        )
        runtime = util_models.RuntimeOptions()
        try:
            rst = await self.client.create_certificate_for_package_request_with_options_async(
                create_certificate_for_package_request_request, runtime
            )
            print(f"域名 {domain} 证书创建成功，订单号: {rst.body.order_id}")
            return {rst.body.order_id: domain}
        except Exception as error:
            print("创建证书时出错:", error.message)
            return None

    async def __delete_cert(self, cert_id: int) -> bool:
        """删除证书"""
        delete_user_certificate_request = (
            cas_20200407_models.DeleteUserCertificateRequest(cert_id=cert_id)
        )
        runtime = util_models.RuntimeOptions()
        try:
            await self.client.delete_user_certificate_with_options_async(
                delete_user_certificate_request, runtime
            )
            return True
        except Exception as error:
            print("删除证书出错：", error.message)
            return False

    async def __verify_cert(self, order_id: int, domain: str) -> dict:
        """
        验证单个证书状态，获取证书信息
        return: dict{order_id,type,certificate,privateKey,domain}
        """
        describe_certificate_state_request = (
            cas_20200407_models.DescribeCertificateStateRequest(order_id=order_id)
        )
        runtime = util_models.RuntimeOptions()
        try:
            rst = await self.client.describe_certificate_state_with_options_async(
                describe_certificate_state_request, runtime
            )
            return {
                "order_id": order_id,
                "type": rst.body.type,
                "certificate": rst.body.certificate,
                "privateKey": rst.body.private_key,
                "domain": domain,
            }
        except Exception as error:
            print("验证证书状态出错:", error.message)

    def __delete_expire_certs(self):
        """删除过期证书"""
        print("删除过期证书...")
        # 获取过期证书列表
        cert_list = self.__get_cert_list(status="EXPIRED", order_type="CERT")
        delete_tasks = []
        # 遍历创建任务
        for cert in cert_list:
            domain = cert.domain
            # 跳过非指定产品和指定域名证书
            if domain not in self.domain_list:
                continue
            delete_tasks.append(self.__delete_cert(cert.certificate_id))

        # 删除证书
        if delete_tasks.__len__() > 0:
            # 异步操作
            async def delete():
                await asyncio.gather(*delete_tasks)

            asyncio.run(delete())

    def __renewal_certs(self) -> dict:
        """
        续签证书
        :return: 需要监控的订单列表 {order_id: domain}
        """
        # 获取当前时间
        now = datetime.datetime.now().timestamp()

        # 获取所有证书列表
        print("检查证书状态...")
        cert_list = self.__get_cert_list()
        renewal_list = dict.fromkeys(self.domain_list, {})  # 不用申请的证书列表
        checking_list = {}  # 审核中的证书列表

        # 遍历证书列表
        for cert in cert_list:
            domain = cert.domain
            # 跳过非指定产品和指定域名证书
            if (
                cert.product_code != config.PRODUCT_CODE
                or domain not in self.domain_list
            ):
                continue

            status = cert.status
            end_time = (cert.cert_end_time or -math.inf) / 1000
            order_id = cert.order_id

            # 证书状态为审核中的直接进行下一步
            if status in ["CHECKING", "PAYED"]:
                renewal_list.pop(domain)
                checking_list[order_id] = domain
                continue

            # 证书到期时间大于阈值，从计划中移除
            if (
                end_time - now > config.EXPIRE_THRESHOLD
                and renewal_list.get(domain) != None
            ):
                print(
                    "证书 %s 剩余时间 %d 天，无需更新"
                    % (domain, (end_time - now) / 86400)
                )
                renewal_list.pop(domain)
                checking_list[order_id] = domain
                continue

            # 证书即将到期, 加入计划列表
            domain_from_list = renewal_list.get(domain)
            if (
                domain_from_list != None
                and domain_from_list.get("end_time", -math.inf) < end_time
            ):
                renewal_list[domain] = {
                    "end_time": end_time,
                    "order_id": order_id,
                    "status": status,
                }

        # 更新证书
        async def update():
            return await asyncio.gather(
                *[self.__create_cert(domain) for domain in renewal_list]
            )

        rst: list = asyncio.run(update())
        for i in rst:
            checking_list.update(i)
        return checking_list

    async def __check_order_status(self, order_info: dict) -> None:
        """
        批量检查证书状态
        :param order_info: [{order_id: domain}]
        :return: 还需继续检查的订单 {order_id: domain}
        """
        rst = await asyncio.gather(
            *[
                self.__verify_cert(order_id, domain)
                for order_id, domain in order_info.items()
            ]
        )

        # 遍历判断订单状态
        for r in rst:
            # 证书更新成功
            if r["type"] == "certificate":
                self.finish_list[r["domain"]] = {
                    "certificate": r["certificate"],
                    "privateKey": r["privateKey"],
                }
                order_info.pop(r["order_id"])

        # 如果有订单未完成，则继续检查
        uncompleted_num = order_info.__len__()
        print("剩余未完成订单数量:", uncompleted_num)
        if uncompleted_num > 0:
            for i in range(config.DELAY):
                print(
                    "等待下次检测，剩余时间: {:02d}".format(config.DELAY - i - 1),
                    end="\r",
                    flush=True,
                )
                await asyncio.sleep(1)
            print()
            await self.__check_order_status(order_info)
            return
        print("证书申请完成...")

    async def __save_cert(self):
        """保存证书到文件"""

        async def check(fn, content):
            if await tool.read(fn) == content:
                return
            await tool.save(fn, content)

        tasks = []
        # 遍历证书列表，判断是否需要保存
        for domain, cert_info in self.finish_list.items():
            tasks.append(check(f"{domain}.pem", cert_info["certificate"]))
            tasks.append(check(f"{domain}.key", cert_info["privateKey"]))

        await asyncio.gather(*tasks)

    def __get_cert_name(self) -> dict:
        """
        获取域名对应证书名称
        return {domain:{end_time, name}}
        """
        rst = self.__get_cert_list(order_type="CERT")
        cache_list = dict.fromkeys(self.domain_list)
        for r in rst:
            domain = r.common_name
            data_in_list = cache_list.get(domain)
            if data_in_list != None and data_in_list["end_date"] > r.end_date:
                continue
            cache_list[domain] = {"end_date": r.end_date, "name": r.name}
        return cache_list

    def run(self) -> dict:
        """
        主任务
        return: 证书信息
        """
        # 查询证书余额
        balance = self.__get_cert_balance()
        print("证书余额为:", balance)
        if balance <= 0:
            print("证书余额不足，请充值！！！")
            return
        # 删除所有过期证书
        self.__delete_expire_certs()
        # 续签证书
        checking_list = self.__renewal_certs()
        # 检查订单状态
        print("正在检查订单状态...")
        asyncio.run(self.__check_order_status(checking_list))
        # 保存证书
        print("正在写入文件...")
        asyncio.run(self.__save_cert())
        print("本地文件写入完毕")
        # 获取域名最新证书的名称
        rst = self.__get_cert_name()
        for k, v in rst.items():
            if k not in self.domain_list:
                continue
            self.finish_list[k]["cert_name"] = v["name"]
        return self.finish_list
